
from typing import Annotated
from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import StreamingResponse

from config.LLMConfig import MODEL_NAME
from graph_agent.charts_agent.agent import create_charts_react_agent
from graph_agent.react_agent.agent import create_custom_react_agent
from models.chat_history_model import ChatHistoryCreate, ChatHistoryResponse
from models.chat_session_model import ChatSessionParams
from models.chat_model import ChatParams
from core.langchain_retrieval import build_history_template, graph_stream, chat_history_crud, build_qa_chain, \
    generate_stream, react_agent_stream
from routers.base import success


router = APIRouter(
    prefix="/chat",
    tags=["chat"],
    responses={404: {"message": "您所访问的资源不存在！"}},
)



@router.post("")
async def chatting(data: ChatParams):
    if not data.messages:
        raise HTTPException(status_code=500, detail="网络异常，请稍后重试！")

    # 先获取历史记录
    history_list = chat_history_crud.list_by_chat_session_id(data.chat_session_id)
    # 再保存 user 消息到历史记录中
    user_chat = ChatHistoryCreate(
        role=data.messages.role,
        content=data.messages.content,
        chat_session_id=data.chat_session_id,
    )
    chat_history_crud.add_item(user_chat)

    # 历史记录转换成LangChain提示词模板
    history_message = build_history_template(history_list)
    # LangChain 检索链 astream() 的参数
    invoke_params = {"question": data.messages.content, "chat_history": history_message}

    try:
        if getattr(data, 'agent', None)=="rag":
            chain = build_qa_chain()

            return StreamingResponse(
                # generate_stream(chain,MODEL_NAME, invoke_params, data.chat_session_id),
                generate_stream(chain, MODEL_NAME, invoke_params, data.chat_session_id),
                media_type="application/x-ndjson",
            )
        elif getattr(data, 'agent', None)=="agent":
            chain = await create_custom_react_agent()

            return StreamingResponse(
                # generate_stream(chain,MODEL_NAME, invoke_params, data.chat_session_id),
                react_agent_stream(chain, MODEL_NAME, invoke_params, data.chat_session_id),
                media_type="application/x-ndjson",
            )
        elif getattr(data, 'agent', None)=="charts":
            chain = await create_charts_react_agent()

            return StreamingResponse(
                # generate_stream(chain,MODEL_NAME, invoke_params, data.chat_session_id),
                react_agent_stream(chain, MODEL_NAME, invoke_params, data.chat_session_id),
                media_type="application/x-ndjson",
            )
        else:
            return StreamingResponse(
                # generate_stream(chain,MODEL_NAME, invoke_params, data.chat_session_id),
                "{}",
                media_type="application/x-ndjson",
            )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"流式响应失败：{str(e)}")


# chat 返回响应流

# graph 返回响应流

@router.get("/history", response_model=ChatHistoryResponse)
async def chat_history(params: Annotated[ChatSessionParams, Query()]):
    results = chat_history_crud.list_by_chat_session_id(params.id)
    return success(results)
